/*
 *  Copyright 2008-2011 NVIDIA Corporation
 *
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */


#include <cusp/detail/config.h>
#include <cusp/array1d.h>

namespace cusp
{
  namespace detail
  {

    template<typename ValueType>
    struct expression_wrapper;

    template<typename LeftParameter, typename RightParameter>
    struct parameter_pair {    
      LeftParameter left;
      RightParameter right;
      typedef typename LeftParameter::ValueType ValueType;
      parameter_pair(const LeftParameter &l, const RightParameter &r) : left(l), right(r) { }     
    };
      
    template<class LeftOperator, class RightOperator, class LeftParameter, class RightParameter>
    struct addition_operator{
      typedef typename LeftParameter::ValueType ValueType;
      typedef parameter_pair<LeftParameter, RightParameter> parameters;
      __host__ __device__ static ValueType exec(int i, const parameters &p) {
	return LeftOperator::exec(i,p.left) + RightOperator::exec(i,p.right); 
      }
    };
      
    template<class LeftOperator, class RightOperator, class LeftParameter, class RightParameter>
    struct multiplication_operator{
      typedef typename LeftParameter::ValueType ValueType;
      typedef parameter_pair<LeftParameter, RightParameter> parameters;
      __host__ __device__ static ValueType exec(int i, const parameters &p) {
	return LeftOperator::exec(i,p.left) * RightOperator::exec(i,p.right); 
      }
    };
      
      
    template<typename T>
    struct array_parameter {
      typedef typename T::value_type ValueType;
      typedef typename T::memory_space memory_space;
      const ValueType * p;
      __host__ __device__ ValueType value(int i) const { 
	return p[i];
      };
	
      array_parameter(typename expression_wrapper<T>::array_type &array,cusp::host_memory){
	p = thrust::raw_pointer_cast(&array[0]);
      }    
      array_parameter(typename expression_wrapper<T>::array_type &array,cusp::device_memory){
	p = thrust::raw_pointer_cast(&array[0]);
      }    
    };


    template<typename T>
    struct constant_parameter{
      typedef T ValueType;
      ValueType v;
      __host__ __device__ ValueType value(int ) const { 
	return v;
      };    
      constant_parameter(typename expression_wrapper<T>::constant_type a):v(a){}
    };
      
    template<class Parameter>
    struct leaf_operator {
      typedef typename Parameter::ValueType ValueType;
      __host__ __device__ static ValueType exec(int i,  const Parameter &p) {	
	return p.value(i);
      }
    };


    template<>
    struct expression_wrapper<float>{
      typedef float constant_type;
      typedef float valid_type;      
      typedef leaf_operator<constant_parameter<float> > wrapped_operator;
      typedef constant_parameter<float> wrapped_parameter;

      template<typename T>
      static operator_parameter_pair<leaf_operator<constant_parameter<T> >,
				     constant_parameter<T> >
      wrap(T a){
	return operator_parameter_pair<leaf_operator<constant_parameter<T> >, 
				       constant_parameter<T> >(constant_parameter<T>(a));
      }
    };

    template<>
    struct expression_wrapper<const float>{
      typedef const float constant_type;
      typedef const float valid_type;      
      typedef leaf_operator<constant_parameter<const float> > wrapped_operator;
      typedef constant_parameter<const float> wrapped_parameter;
      
      template<typename T>
      static operator_parameter_pair<leaf_operator<constant_parameter<const T> >,
				     constant_parameter<const T> >
      wrap(const T a){
	return operator_parameter_pair<leaf_operator<constant_parameter<const T> >, 
				       constant_parameter<const T> >(constant_parameter<const T>(a));
      }      
    };

    template<typename T, typename M>
    struct expression_wrapper<cusp::array1d<T,M> >{
      typedef cusp::array1d<T,M> array_type;
      typedef cusp::array1d<T,M> valid_type;      
      typedef leaf_operator<array_parameter<cusp::array1d<T,M> > > wrapped_operator;
      typedef array_parameter<cusp::array1d<T,M> > wrapped_parameter;
      
      static operator_parameter_pair<leaf_operator<array_parameter<const cusp::array1d<T,M> > >,
				     array_parameter<const cusp::array1d<T,M> > >
      wrap(cusp::array1d<T,M> & a){
	return operator_parameter_pair<leaf_operator<array_parameter<cusp::array1d<T,M> > >, 
				       array_parameter<cusp::array1d<T,M> > >(array_parameter<cusp::array1d<T,M> > (a,M()));
      }	  
    };

    template<typename T, typename M>
    struct expression_wrapper<const cusp::array1d<T,M> >{
      typedef const cusp::array1d<T,M> array_type;
      typedef const cusp::array1d<T,M> valid_type;      
      typedef leaf_operator<array_parameter<const cusp::array1d<T,M> > > wrapped_operator;
      typedef array_parameter<const cusp::array1d<T,M> > wrapped_parameter;
      
      static operator_parameter_pair<leaf_operator<array_parameter<const cusp::array1d<T,M> > >,
				     array_parameter<const cusp::array1d<T,M> > >
      wrap(const cusp::array1d<T,M> & a){
	return operator_parameter_pair<leaf_operator<array_parameter<const cusp::array1d<T,M> > >, 
				       array_parameter<const cusp::array1d<T,M> > >(array_parameter<const cusp::array1d<T,M> > (a,M()));
      }
    };

    template<typename Operator, typename Parameter>
    struct expression_wrapper<operator_parameter_pair<Operator,Parameter> >{
      typedef operator_parameter_pair<Operator,Parameter> valid_type;      
      typedef Operator wrapped_operator;
      typedef Parameter wrapped_parameter;
      
      static operator_parameter_pair<Operator,Parameter> wrap(operator_parameter_pair<Operator,Parameter> a){
	return a;
      }
    };

    template<typename Operator, typename Parameter>
    struct expression_wrapper<const operator_parameter_pair<Operator,Parameter> >{
      typedef const operator_parameter_pair<Operator,Parameter> valid_type;      
      typedef Operator wrapped_operator;
      typedef Parameter wrapped_parameter;
      static operator_parameter_pair<Operator,Parameter> wrap(const operator_parameter_pair<Operator,Parameter> a){
	return a;
      }
      
    };  


    template <typename Expression, typename T>
    struct fused_operator: public thrust::unary_function<int,T>
    {
      const Expression func;
      fused_operator(const Expression _func) : func(_func) {}
      __host__ __device__ T operator()(int i){
	return func.exec(i);
      }
    };

    template<class Operator, class Parameter>
    struct operator_parameter_pair {     
      Parameter parm;
      typedef typename Parameter::ValueType ValueType;
      operator_parameter_pair(const Parameter &p) : parm(p) { }    
      __host__ __device__ 
      ValueType exec(int i) const {
	return Operator::exec(i, parm); 
      }
    };

    
    template<class LeftOperator, class LeftParameter, class RightOperator, class RightParameter>
    operator_parameter_pair<addition_operator<LeftOperator, 
					      RightOperator,
					      LeftParameter, 
					      RightParameter>, 
			    parameter_pair<LeftParameter, 
					   RightParameter> >
    operator+(const operator_parameter_pair<LeftOperator, LeftParameter> &left,
	      const operator_parameter_pair<RightOperator, RightParameter> &right){
      return operator_parameter_pair<addition_operator<LeftOperator, RightOperator,
						       LeftParameter, RightParameter>, 
				     parameter_pair<LeftParameter, RightParameter> >
	(parameter_pair<LeftParameter, RightParameter>(left.parm, right.parm));
    };
      
    template<class LeftOperator, class LeftParameter, class RightOperator, class RightParameter>
    operator_parameter_pair<multiplication_operator<LeftOperator, 
						    RightOperator,
						    LeftParameter, 
						    RightParameter>, 
			    parameter_pair<LeftParameter, 
					   RightParameter> >
    operator*(const operator_parameter_pair<LeftOperator, LeftParameter> &left,
	      const operator_parameter_pair<RightOperator, RightParameter> &right){
      return operator_parameter_pair<multiplication_operator<LeftOperator, RightOperator,
							     LeftParameter, RightParameter>, 
				     parameter_pair<LeftParameter, RightParameter> >
	(parameter_pair<LeftParameter, RightParameter>(left.parm, right.parm));
    };
  }
}


namespace cusp
{  
  template <typename Expression, typename Array>
  void fused_execute(const Expression &func, Array &dst){
    thrust::counting_iterator<int> first(0);
    thrust::counting_iterator<int> last = first + dst.size();
    thrust::transform(first,last,dst.begin(),detail::fused_operator<Expression,typename Array::value_type>(func));
  }      
  
  
  template <typename LeftType, typename RightType>
  cusp::detail::operator_parameter_pair<cusp::detail::addition_operator<typename cusp::detail::expression_wrapper<const LeftType>::wrapped_operator,
									typename cusp::detail::expression_wrapper<const RightType>::wrapped_operator,
									typename cusp::detail::expression_wrapper<const LeftType>::wrapped_parameter,
									typename cusp::detail::expression_wrapper<const RightType>::wrapped_parameter>,
					cusp::detail::parameter_pair<typename cusp::detail::expression_wrapper<const LeftType>::wrapped_parameter,
								     typename cusp::detail::expression_wrapper<const RightType>::wrapped_parameter> >
  operator+(const LeftType & a, const RightType & b){
    return cusp::detail::expression_wrapper<const LeftType>::wrap(a) + cusp::detail::expression_wrapper<const RightType>::wrap(b);
  }    
}
  
